package com.sinosoft.common.encrypt.filter;

import com.sinosoft.common.core.enums.ErrorCodeEnum;
import com.sinosoft.common.core.exception.GlobalAuthException;
import com.sinosoft.common.core.utils.StringUtils;
import com.sinosoft.common.encrypt.properties.ApiDecryptProperties;
import com.sinosoft.common.encrypt.utils.CryptoUtils;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import lombok.extern.slf4j.Slf4j;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;

/**
 * JWT安全模式响应参数包装类
 * 使用SM2/SM4加密方式 (SAFE_MODE=1)
 *
 * @author zzf
 */
@Slf4j
public class JwtSafeModeResponseBodyWrapper extends HttpServletResponseWrapper {

    private final ByteArrayOutputStream buffer;
    private final ServletOutputStream output;
    private final PrintWriter writer;

    public JwtSafeModeResponseBodyWrapper(HttpServletResponse response) throws IOException {
        super(response);
        this.buffer = new ByteArrayOutputStream();
        this.output = createOutputStream();
        this.writer = new PrintWriter(new OutputStreamWriter(buffer, StandardCharsets.UTF_8), true);
    }

    @Override
    public PrintWriter getWriter() {
        return writer;
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return output;
    }

    @Override
    public void flushBuffer() throws IOException {
        if (writer != null) {
            writer.flush();
        }
        if (output != null) {
            output.flush();
        }
    }

    @Override
    public void reset() {
        buffer.reset();
    }

    /**
     * 获取响应数据字节数组
     */
    public byte[] getResponseData() throws IOException {
        flushBuffer();
        return buffer.toByteArray();
    }

    /**
     * 获取响应内容字符串
     */
    public String getContent() throws IOException {
        flushBuffer();
        return buffer.toString(StandardCharsets.UTF_8);
    }

    /**
     * 获取加密内容
     *
     * @param servletResponse response
     * @param messageId       消息ID
     * @param encryptKey      加密密钥
     * @param properties      配置属性
     * @return 加密后的响应内容
     * @throws IOException IO异常
     */
    public String getEncryptContent(HttpServletResponse servletResponse, String messageId, String encryptKey, ApiDecryptProperties properties) throws IOException {
        // 验证SM2私钥配置
        if (StringUtils.isBlank(properties.getSm2PrivateKey())) {
            throw new GlobalAuthException(ErrorCodeEnum.PARAM_EMPTY, "SM2私钥未配置", "JWT安全模式需要配置SM2私钥");
        }
        servletResponse.setCharacterEncoding(StandardCharsets.UTF_8.toString());
        String originalBody = this.getContent();
        // JWT安全模式保持原有响应格式，只改变加密方式
        // 直接对原始响应内容进行SM4加密
        return encryptResponseIfNeeded(originalBody, encryptKey, properties);
    }

    /**
     * 创建自定义ServletOutputStream
     */
    private ServletOutputStream createOutputStream() {
        return new ServletOutputStream() {
            @Override
            public boolean isReady() {
                return true;
            }

            @Override
            public void setWriteListener(WriteListener writeListener) {
                // 不需要实现
            }

            @Override
            public void write(int b) throws IOException {
                buffer.write(b);
            }

            @Override
            public void write(byte[] b) throws IOException {
                buffer.write(b);
            }

            @Override
            public void write(byte[] b, int off, int len) throws IOException {
                buffer.write(b, off, len);
            }
        };
    }

    /**
     * 根据需要加密响应内容
     */
    private String encryptResponseIfNeeded(String content, String encryptKey, ApiDecryptProperties properties) {
        if (StringUtils.isNotBlank(encryptKey)) {
            try {
                // 解密SM4密钥
                String sm4Key = CryptoUtils.decryptSm4Key(encryptKey, properties.getSm2PrivateKey());
                // 加密响应内容
                return CryptoUtils.encryptBody(content, sm4Key);
            } catch (Exception e) {
                log.error("JWT安全模式响应加密失败: {}", e.getMessage());
                // 加密失败时抛出异常，不返回明文
                throw new GlobalAuthException(ErrorCodeEnum.SYSTEM_INTERNAL_ERROR, "响应加密失败", e.getMessage());
            }
        }
        return content;
    }
}
